import os
import requests
import pandas as pd
from tqdm import tqdm
from multiprocessing import Pool
from requests.adapters import HTTPAdapter

s = requests.Session()
s.mount('http://', HTTPAdapter(max_retries=1))
s.mount('https://', HTTPAdapter(max_retries=1))

class Downloader(object):
    def __init__(self, save_root):
        self.save_root = save_root
    
    @staticmethod
    def __download(files):
        file_path, url = files
        chunk_size = 1024
        response = requests.get(url, stream=True, timeout=3)
        if response.status_code == 200:
            file_size = int(response.headers['content-length'])
            file_name = os.path.basename(file_path)
            pbar = tqdm(range(file_size), colour='green', unit='B', desc='download %s' % file_name, unit_scale=True, unit_divisor=chunk_size, leave=False)

            with open(file_path, 'ab') as out_file:
                for data in response.iter_content(chunk_size=chunk_size):
                    if data:
                        out_file.write(data)
                    pbar.update()
    
    def download(self, files, nproc=1):
        # multi-process download files
        with Pool(nproc) as p:
            tqdm(p.imap(self.__download, files), total=len(files), colour='blue', leave=True)
    
    def load_list(self, list_path, ext=''):
        files = []
        with open(list_path, 'r') as f:
            for line in f.readlines():
                name, url = line.strip().split('\t')
                name = name + ext
                files.append((name, url))
        return files
    
    def load_xlsx(self, xlsx_path, url_head, name_head, ext=''):
        files = []
        if xlsx_path.endswith('.xlsx'):
            df = pd.read_excel(xlsx_path)
        elif xlsx_path.endswith('.csv'):
            df = pd.read_csv(xlsx_path)
        else:
            raise ValueError('Unsupported file type: %s' % xlsx_path)

        for i in range(df.shape[0]):
            values = df.loc[i]
            name = values[name_head] + ext
            url = values[url_head]
            files.append((name, url))
        return files
        


           